import torch.optim
import torchvision.models as models
import torch.nn as nn  # 这里将重要的库命名。
import torchvision.transforms as transforms
from torchvision import datasets
from torch.utils.data import DataLoader
import gc
import os
from PIL import Image  # 这里直接导入3个经典的包。
from torch.utils.data import Dataset

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class Re_net(nn.Module):
    def __init__(self, classes):  # 这里的classes表示的是总共有多少个分类。
        super(Re_net, self).__init__()
        self.net = models.resnet34(
            weights=models.ResNet34_Weights.IMAGENET1K_V1)  # 使用已经定义好的模型来训练。weights=models.ResNet34_Weights.IMAGENET1K_V1
        self.num1 = self.net.fc.in_features  # 将最后的那个输出提取出来。
        self.net.fc = nn.Linear(self.num1, classes)

    def forward(self, x):  # 这里的self表示的是输入的数据。
        x = self.net(x)
        return x


class MyData(Dataset):  # 这里的filenames表示的是对应图片的文件夹，也就是分类。
    def __init__(self, filenames):
        self.filenames = filenames
        self.path = os.listdir(self.filenames)  # 这里提取出对应图片的名称。

    def __getitem__(self, idx):
        im_name = self.path[idx]
        im_path = os.path.join(self.filenames, im_name)
        image = Image.open(im_path)
        return image

    def __len__(self):
        return len(self.path)

    def get_name(self):
        return self.filenames


def tr_te():
    torch.cuda.manual_seed_all(2233)
    trans = transforms.Compose(  # 这里的转化指的是首先将数据集转化为224*224，然后转化为tensor,最后再标准化。
        [transforms.Resize([224, 224]), transforms.RandomHorizontalFlip(p=0.5), transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])  # 这里是分类任务中常常使用的标准化。
    train_data = datasets.ImageFolder("./dataset/train", transform=trans)
    test_data = datasets.ImageFolder("./dataset/test", transform=trans)
    train = DataLoader(train_data, batch_size=32, shuffle=True, drop_last=False)  # 每次划分32个即可。
    test = DataLoader(test_data, batch_size=32, shuffle=False, drop_last=False)
    return train, test  # 这里直接返回打包好的数据集。


def train_function(net, epoch, optimizer, loss_function, train, test):  # 传进来的参数分别是网络，迭代次数，优化方法，损失函数，训练集，测试集
    loss_list = []  # 创建一个空的列表，用于储存损失值。
    train_acclist = []  # 这里也是一个空的列表。用于储存训练集的准确率。
    test_acclist = []  # 也用来储存测试集的准确率。
    for i in range(epoch):
        loss = 0  # 初始化loss值。
        train_acc = 0
        test_acc = 0

        net.train()  # 这里由于使用了normalize,所以最好将网络弄出train模式。
        for images, labels in train:
            images = images.to(device)  # 这里直接将数据转化到GPU上面。
            labels = labels.to(device)
            optimizer.zero_grad()
            y_pre = net(images)  # 这里直接将数据
            loss_single = loss_function(y_pre, labels)  # 这里计算本次batch_size的损失值的平均值。
            # 为了防止显存爆炸，不能直接相加。
            loss += loss_single.item() * images.size(0)  # 乘以总的batch数量就是本次的损失值。
            loss_single.backward()  # 进行反向传播即可。
            optimizer.step()

            _, pre_data_train = torch.max(y_pre.detach(), 1)  # 求出来此时训练集的准确率
            train_acc += torch.mean((pre_data_train == labels).float())  # 这里可以将bool转化为浮点数然后求和。
        # print(train_acc)

        net.eval()  # 这里是测试数据，为了防止显存爆炸可以使用这种方法。
        for images, labels in test:
            images = images.to(device)
            labels = labels.to(device)  # 弄到GPU上面即可。
            test_pre = net(images)  # 使用已经训练好的网络来
            # 这里测试集就不计算损失值了。
            _, pre_data_train = torch.max(test_pre.detach(), 1)
            test_acc += torch.mean((pre_data_train == labels).float())
        # 接下来每次的enpch都输出这些数字即可。
        train_loss = loss / len(train.dataset)
        train_acc /= len(train)
        test_acc /= len(test)
        # 这里将每次epoch的值储存起来即可。
        train_acclist.append(train_acc)
        loss_list.append(train_loss)
        test_acclist.append(test_acc)

        print(str(i + 1) + " 损失值为：" + str(float(train_loss)) + ".训练集的准确率为:" + str(
            float(train_acc) * 100) + "%" + "测试集的准确率为：" + str(float(test_acc) * 100) + "%")
        gc.collect()
        torch.cuda.empty_cache()  # 每一次epoch都进行显存释放。
    return net, train_acclist, test_acclist, loss_list  # 这里一定要返回已经训练好的net.


def test_function(net, test):  # 这里传入的参数是已经训练好的网络以及测试集。
    net.eval()
    acc = 0
    for images, labels in test:
        images = images.to(device)
        labels = labels.to(device)  # 转化到显卡上面即可。
        finish = net(images)
        _, pre = torch.max(finish.detach(), 1)
        acc += torch.sum((pre == labels).float())
    print("测试集的准确率为：" + str(float(100 * acc / 1436)) + "%")


# 这里的数据增强主要是对数据进行翻转.
def v_function(image):
    v = transforms.RandomHorizontalFlip(p=1)  # 翻转，水平。
    image = v(image)
    return image


def h_function(image):  # 这里是将数据进行垂直翻转。
    h = transforms.RandomVerticalFlip(p=1)
    image = h(image)
    return image
